package com.zengde.ai_sandbox.service;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ArrayNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import jakarta.annotation.PreDestroy;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.function.client.WebClientRequestException;
import org.springframework.web.reactive.function.client.WebClientResponseException;
import reactor.core.publisher.Mono;
import reactor.util.retry.Retry;

import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Qdrant向量数据库服务
 * 提供与Qdrant数据库的基本交互功能
 * 使用HTTP API而不是Java客户端，以避免API兼容性问题
 */
@Service
@Slf4j
public class QdrantService {

    private final WebClient webClient;
    private final ObjectMapper objectMapper;
    private static final String COLLECTION_NAME = "demo_collection1";
    private static final int VECTOR_SIZE = 4; // 使用简单的4维向量作为示例
    private static final int MAX_RETRIES = 3;
    private static final Duration RETRY_DELAY = Duration.ofSeconds(1);
    private boolean qdrantAvailable = false;

    public QdrantService() {
        // 创建WebClient连接到Qdrant服务器
        this.webClient = WebClient.builder()
                .baseUrl("http://localhost:6333")
                .defaultHeaders(headers -> {
                    headers.setContentType(MediaType.APPLICATION_JSON);
                    // 移除基本认证，除非Qdrant服务器有设置
                })
                .build();
        
        this.objectMapper = new ObjectMapper();
        
        // 检查Qdrant服务是否可用
        checkQdrantAvailability();
    }
    
    /**
     * 检查Qdrant服务是否可用
     */
    private void checkQdrantAvailability() {
        try {
            String response = webClient.get()
                    .uri("/")
                    .retrieve()
                    .bodyToMono(String.class)
                    .retryWhen(Retry.fixedDelay(1, Duration.ofSeconds(1))
                            .filter(e -> e instanceof WebClientRequestException))
                    .onErrorReturn("{\"status\":\"unavailable\"}")
                    .block();
            
            qdrantAvailable = response != null && !response.contains("unavailable");
            
            if (qdrantAvailable) {
                log.info("Qdrant服务已初始化，连接到http://localhost:6334");
            } else {
                log.warn("Qdrant服务不可用，请确保Qdrant已启动在http://localhost:6334");
            }
        } catch (Exception e) {
            log.error("检查Qdrant服务可用性时出错: {}", e.getMessage(), e);
            qdrantAvailable = false;
        }
    }

    /**
     * 检查集合是否存在
     */
    public boolean collectionExists(String collectionName) {
        if (!qdrantAvailable) {
            log.warn("Qdrant服务不可用，无法检查集合是否存在");
            return false;
        }
        
        try {
            // 获取所有集合
            String response = webClient.get()
                    .uri("/collections")
                    .retrieve()
                    .bodyToMono(String.class)
                    .retryWhen(Retry.fixedDelay(MAX_RETRIES, RETRY_DELAY)
                            .filter(e -> e instanceof WebClientRequestException))
                    .onErrorResume(e -> {
                        log.error("获取集合列表时出错: {}", e.getMessage());
                        return Mono.just("{}");
                    })
                    .block();
            
            if (response == null || response.equals("{}")) {
                return false;
            }
            
            JsonNode responseJson = objectMapper.readTree(response);
            JsonNode collectionsNode = responseJson.path("result").path("collections");
            
            // 检查指定名称的集合是否存在
            if (collectionsNode.isArray()) {
                for (JsonNode collection : collectionsNode) {
                    if (collection.path("name").asText().equals(collectionName)) {
                        return true;
                    }
                }
            }
            return false;
        } catch (Exception e) {
            log.error("检查集合存在时出错: {}", e.getMessage(), e);
            return false;
        }
    }

    /**
     * 创建一个演示集合
     */
    public void createCollection() {
        if (!qdrantAvailable) {
            throw new RuntimeException("Qdrant服务不可用，请确保Qdrant已启动在http://localhost:6334");
        }
        
        try {
            // 检查集合是否已存在
            if (collectionExists(COLLECTION_NAME)) {
                log.info("集合已存在: {}", COLLECTION_NAME);
                return;
            }
            
            // 创建集合请求体
            ObjectNode requestBody = objectMapper.createObjectNode();
            ObjectNode vectorsConfig = requestBody.putObject("vectors");
            vectorsConfig.put("size", VECTOR_SIZE);
            vectorsConfig.put("distance", "Cosine");
            
            // 发送创建集合请求
            String response = webClient.put()
                    .uri("/collections/{collection_name}", COLLECTION_NAME)
                    .bodyValue(requestBody)
                    .retrieve()
                    .bodyToMono(String.class)
                    .retryWhen(Retry.fixedDelay(MAX_RETRIES, RETRY_DELAY)
                            .filter(e -> e instanceof WebClientRequestException))
                    .onErrorResume(e -> {
                        if (e instanceof WebClientResponseException) {
                            WebClientResponseException wcre = (WebClientResponseException) e;
                            log.error("创建集合失败，状态码: {}, 响应: {}", 
                                     wcre.getStatusCode(), wcre.getResponseBodyAsString());
                            return Mono.just("创建集合失败: " + wcre.getStatusCode());
                        } else {
                            log.error("创建集合时出错: {}", e.getMessage(), e);
                            return Mono.just("创建集合失败: " + e.getMessage());
                        }
                    })
                    .block();
            
            log.info("已成功创建集合: {}, 响应: {}", COLLECTION_NAME, response);
        } catch (Exception e) {
            log.error("创建集合时出错: {}", e.getMessage(), e);
            throw new RuntimeException("创建集合失败: " + e.getMessage(), e);
        }
    }

    /**
     * 添加单个向量
     */
    public void addVector(Long id, List<Float> vectorData) {
        if (!qdrantAvailable) {
            throw new RuntimeException("Qdrant服务不可用，请确保Qdrant已启动在http://localhost:6334");
        }
        
        try {
            // 验证向量尺寸
            if (vectorData.size() != VECTOR_SIZE) {
                throw new IllegalArgumentException(
                    String.format("向量必须是%d维的，但提供了%d维", VECTOR_SIZE, vectorData.size()));
            }
            
            // 创建请求体
            ObjectNode requestBody = objectMapper.createObjectNode();
            ArrayNode points = requestBody.putArray("points");
            
            // 创建点
            ObjectNode point = points.addObject();
            point.put("id", id);
            
            // 添加向量数据
            ArrayNode vector = point.putArray("vector");
            for (Float value : vectorData) {
                vector.add(value);
            }
            
            // 添加有效载荷
            ObjectNode payload = point.putObject("payload");
            payload.put("id", id);
            
            // 发送添加向量请求
            String response = webClient.put()
                    .uri("/collections/{collection_name}/points", COLLECTION_NAME)
                    .bodyValue(requestBody)
                    .retrieve()
                    .bodyToMono(String.class)
                    .retryWhen(Retry.fixedDelay(MAX_RETRIES, RETRY_DELAY)
                            .filter(e -> e instanceof WebClientRequestException))
                    .onErrorResume(e -> {
                        if (e instanceof WebClientResponseException) {
                            WebClientResponseException wcre = (WebClientResponseException) e;
                            log.error("添加向量失败，状态码: {}, 响应: {}", 
                                     wcre.getStatusCode(), wcre.getResponseBodyAsString());
                            return Mono.just("添加向量失败: " + wcre.getStatusCode());
                        } else {
                            log.error("添加向量时出错: {}", e.getMessage(), e);
                            return Mono.just("添加向量失败: " + e.getMessage());
                        }
                    })
                    .block();
            
            log.info("已成功添加向量: id={}, 响应: {}", id, response);
        } catch (Exception e) {
            log.error("添加向量时出错: {}", e.getMessage(), e);
            throw new RuntimeException("添加向量失败: " + e.getMessage(), e);
        }
    }

    /**
     * 添加示例向量数据到集合
     */
    public void addVectors() {
        if (!qdrantAvailable) {
            throw new RuntimeException("Qdrant服务不可用，请确保Qdrant已启动在http://localhost:6334");
        }
        
        try {
            // 创建示例向量数据
            List<Map<String, Object>> vectorData = createSampleVectorData();
            
            // 创建请求体
            ObjectNode requestBody = objectMapper.createObjectNode();
            ArrayNode points = requestBody.putArray("points");
            
            // 添加所有点
            for (int i = 0; i < vectorData.size(); i++) {
                Map<String, Object> data = vectorData.get(i);
                @SuppressWarnings("unchecked")
                List<Float> vector = (List<Float>) data.get("vector");
                
                // 创建点
                ObjectNode point = points.addObject();
                point.put("id", i + 1);
                
                // 添加向量数据
                ArrayNode vectorNode = point.putArray("vector");
                for (Float value : vector) {
                    vectorNode.add(value);
                }
                
                // 添加有效载荷
                ObjectNode payload = point.putObject("payload");
                payload.put("name", (String) data.get("name"));
                payload.put("description", (String) data.get("description"));
            }
            
            // 发送添加向量请求
            String response = webClient.put()
                    .uri("/collections/{collection_name}/points", COLLECTION_NAME)
                    .bodyValue(requestBody)
                    .retrieve()
                    .bodyToMono(String.class)
                    .retryWhen(Retry.fixedDelay(MAX_RETRIES, RETRY_DELAY)
                            .filter(e -> e instanceof WebClientRequestException))
                    .onErrorResume(e -> {
                        if (e instanceof WebClientResponseException) {
                            WebClientResponseException wcre = (WebClientResponseException) e;
                            log.error("添加向量批量失败，状态码: {}, 响应: {}", 
                                     wcre.getStatusCode(), wcre.getResponseBodyAsString());
                            return Mono.just("添加向量批量失败: " + wcre.getStatusCode());
                        } else {
                            log.error("添加向量批量时出错: {}", e.getMessage(), e);
                            return Mono.just("添加向量批量失败: " + e.getMessage());
                        }
                    })
                    .block();
            
            log.info("已成功添加 {} 个向量到集合, 响应: {}", vectorData.size(), response);
        } catch (Exception e) {
            log.error("添加向量数据时出错: {}", e.getMessage(), e);
            throw new RuntimeException("添加向量失败: " + e.getMessage(), e);
        }
    }

    /**
     * 搜索最相似的向量
     */
    public List<Map<String, Object>> searchSimilar(List<Float> queryVector, int limit) {
        if (!qdrantAvailable) {
            throw new RuntimeException("Qdrant服务不可用，请确保Qdrant已启动在http://localhost:6334");
        }
        
        try {
            // 验证向量尺寸
            if (queryVector.size() != VECTOR_SIZE) {
                throw new IllegalArgumentException(
                    String.format("向量必须是%d维的，但提供了%d维", VECTOR_SIZE, queryVector.size()));
            }
            
            // 创建搜索请求体
            ObjectNode requestBody = objectMapper.createObjectNode();
            ArrayNode vector = requestBody.putArray("vector");
            for (Float value : queryVector) {
                vector.add(value);
            }
            requestBody.put("limit", limit);
            
            // 添加有效载荷选择器
            ObjectNode payloadSelector = requestBody.putObject("with_payload");
            payloadSelector.put("enable", true);
            
            // 发送搜索请求
            String response = webClient.post()
                    .uri("/collections/{collection_name}/points/search", COLLECTION_NAME)
                    .bodyValue(requestBody)
                    .retrieve()
                    .bodyToMono(String.class)
                    .retryWhen(Retry.fixedDelay(MAX_RETRIES, RETRY_DELAY)
                            .filter(e -> e instanceof WebClientRequestException))
                    .onErrorResume(e -> {
                        if (e instanceof WebClientResponseException) {
                            WebClientResponseException wcre = (WebClientResponseException) e;
                            log.error("搜索向量失败，状态码: {}, 响应: {}", 
                                     wcre.getStatusCode(), wcre.getResponseBodyAsString());
                            return Mono.just("{\"result\":[]}");
                        } else {
                            log.error("搜索向量时出错: {}", e.getMessage(), e);
                            return Mono.just("{\"result\":[]}");
                        }
                    })
                    .block();
            
            // 解析响应
            List<Map<String, Object>> results = new ArrayList<>();
            JsonNode responseJson = objectMapper.readTree(response);
            JsonNode resultsNode = responseJson.path("result");
            
            if (resultsNode.isArray()) {
                for (JsonNode result : resultsNode) {
                    Map<String, Object> resultMap = new HashMap<>();
                    resultMap.put("id", result.path("id").asLong());
                    resultMap.put("score", result.path("score").asDouble());
                    
                    // 解析负载
                    Map<String, Object> payload = new HashMap<>();
                    JsonNode payloadNode = result.path("payload");
                    if (payloadNode.isObject()) {
                        payloadNode.fields().forEachRemaining(entry -> {
                            payload.put(entry.getKey(), entry.getValue().asText());
                        });
                    }
                    resultMap.put("payload", payload);
                    
                    results.add(resultMap);
                }
            }
            
            log.info("搜索完成，找到 {} 个结果", results.size());
            return results;
        } catch (Exception e) {
            log.error("搜索相似向量时出错: {}", e.getMessage(), e);
            return List.of();
        }
    }
    
    /**
     * 重载搜索方法，使用默认限制数量5
     */
    public List<Map<String, Object>> searchSimilar(List<Float> queryVector) {
        return searchSimilar(queryVector, 5);
    }

    /**
     * 创建示例向量数据
     */
    private List<Map<String, Object>> createSampleVectorData() {
        List<Map<String, Object>> data = new ArrayList<>();
        
        // 示例1
        Map<String, Object> item1 = new HashMap<>();
        item1.put("vector", List.of(0.1f, 0.2f, 0.3f, 0.4f));
        item1.put("name", "苹果1");
        item1.put("description", "一种常见的水果，红色或绿色，酸甜可口");
        data.add(item1);
        
        // 示例2
        Map<String, Object> item2 = new HashMap<>();
        item2.put("vector", List.of(0.2f, 0.3f, 0.4f, 0.5f));
        item2.put("name", "香蕉1");
        item2.put("description", "一种黄色的水果，弯曲形状，甜味");
        data.add(item2);
        
        // 示例3
        Map<String, Object> item3 = new HashMap<>();
        item3.put("vector", List.of(0.3f, 0.4f, 0.5f, 0.6f));
        item3.put("name", "橙子");
        item3.put("description", "一种橙色的水果，球形，酸甜口味");
        data.add(item3);
        
        // 示例4
        Map<String, Object> item4 = new HashMap<>();
        item4.put("vector", List.of(0.4f, 0.5f, 0.6f, 0.7f));
        item4.put("name", "草莓");
        item4.put("description", "一种红色的浆果，表面有小籽，甜美多汁");
        data.add(item4);
        
        return data;
    }
    
    /**
     * 兼容旧的API调用
     */
    public void createNewCollection() {
        createCollection();
    }
    
    /**
     * 获取Qdrant服务可用性状态
     */
    public boolean isQdrantAvailable() {
        return qdrantAvailable;
    }
}
